#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2019/7/25 下午8:56
# @Author  : fugang_le
# @Software: PyCharm


def get_contextual_word_embedding():
    from bert_serving.client import BertClient
    # max_seq_len = 25
    # pooling_strategy = NONE

    bc = BertClient()
    vec = bc.encode(['hey you', 'whats up?'])

    vec  # [2, 25, 768]
    vec[0]  # [1, 25, 768], sentence embeddings for `hey you`
    vec[0][0]  # [1, 1, 768], word embedding for `[CLS]`
    vec[0][1]  # [1, 1, 768], word embedding for `hey`
    vec[0][2]  # [1, 1, 768], word embedding for `you`
    vec[0][3]  # [1, 1, 768], word embedding for `[SEP]`
    vec[0][4]  # [1, 1, 768], word embedding for padding symbol
    vec[0][25]  # error, out of index!


def build_QA_semantic_search_engine():
    '''
    利用bert获取句向量，然后用余弦相似度计算
    :return:
    '''
    import numpy as np

    service_ip = '172.18.8.35'
    service_port = 8503

    question = '陪产假需要提供什么材料'

    answers = ['流产假需要提供什么材料', '陪产假要交什么材料', '丧假需要提供哪些证明材料', '休病假要交什么材料',
                      '已转深户，我需要提交什么材料？', '如何休陪产假', '哪些休假是需要提交休假材料的？', '休假材料提交到哪？',
                      '陪产假有多少天', '如何提交材料']

    from bert_serving.client import BertClient
    bc = BertClient(ip=service_ip, port=service_port)  # ip address of the GPU machine
    sent_vec = bc.encode([question])[0]
    answers_vecs = bc.encode(answers)
    scores = []
    for index, candidate_vec in enumerate(answers_vecs):
        score = np.dot(sent_vec, candidate_vec) / (np.linalg.norm(sent_vec) * (np.linalg.norm(candidate_vec)))
        scores.append({'sentence': question + "||" + answers[index], "score": score})
    # print(sorted(scores, key=lambda item: item['score'], reverse=True))
    scores.sort(key=lambda x: x['score'], reverse=True)
    print(scores)

    topk = 3
    while True:
        query = input('your question: ')
        query_vec = bc.encode([query])[0]
        # compute normalized dot product as score
        score = np.sum(query_vec * answers_vecs, axis=1) / np.linalg.norm(answers_vecs, axis=1)
        topk_idx = np.argsort(score)[::-1][:topk]
        for idx in topk_idx:
            print('> %s\t%s' % (score[idx], answers[idx]))


if __name__ == '__main__':
    build_QA_semantic_search_engine()